# optimal_knn.py

import pickle
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image  # 新增：用于图像增强的PIL库
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.neighbors import KNeighborsClassifier
from tqdm import tqdm  # 用于显示进度条

# 加载手写数字数据集
digits = load_digits()
X = digits.data
y = digits.target

# 数据增强：轻微旋转、缩放原始图像（模拟不同手写风格）
def augment_image(img):
    img_2d = img.reshape(8, 8)
    # 随机轻微旋转（-10到10度）
    rotated = Image.fromarray(img_2d).rotate(np.random.uniform(-10, 10), expand=False, fillcolor=0)
    # 随机缩放（0.8到1.2倍）- 这里使用旋转后的图像直接作为增强结果
    return np.array(rotated).flatten()

# 对数据集进行增强（扩充数据量）
X_augmented = np.array([augment_image(img) for img in X])
X_combined = np.vstack([X, X_augmented])  # 原始数据+增强数据
y_combined = np.hstack([y, y])  # 标签对应扩充

# 将扩充后的数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_combined, y_combined, test_size=0.2, random_state=42)

# 扩展参数网格：搜索K值、权重方式和距离度量
param_grid = {
    'n_neighbors': range(1, 41),
    'weights': ['uniform', 'distance'],  # 距离加权方式
    'metric': ['euclidean', 'manhattan']  # 距离度量方式
}

# 网格搜索最优参数（替代原有的for循环交叉验证）
print("开始网格搜索最优参数...")
grid_search = GridSearchCV(
    KNeighborsClassifier(),
    param_grid,
    cv=5,
    scoring='accuracy',
    n_jobs=-1  # 并行计算加速
)
grid_search.fit(X_train, y_train)

# 最佳模型和参数
best_knn = grid_search.best_estimator_
best_k = grid_search.best_params_['n_neighbors']
print(f"最优参数: {grid_search.best_params_}")
print(f"最优交叉验证准确率: {grid_search.best_score_:.4f}")

# 计算不同K值在测试集上的准确率（用于可视化）
test_accuracies = []
cv_accuracies = []
k_range = range(1, 41)

# 提取网格搜索结果中不同K值的平均性能
for k in tqdm(k_range, desc="收集不同K值性能数据"):
    # 测试集准确率
    knn = KNeighborsClassifier(
        n_neighbors=k,
        weights=grid_search.best_params_['weights'],
        metric=grid_search.best_params_['metric']
    )
    knn.fit(X_train, y_train)
    test_accuracies.append(knn.score(X_test, y_test))
    
    # 交叉验证平均准确率（从网格搜索结果中提取）
    mask = [param['n_neighbors'] == k for param in grid_search.cv_results_['params']]
    cv_mean = grid_search.cv_results_['mean_test_score'][mask].mean()
    cv_accuracies.append(cv_mean)

# 绘制折线图
plt.figure(figsize=(10, 6))
plt.plot(k_range, test_accuracies, marker='o', linestyle='-', color='blue', label='Test Accuracy')
plt.plot(k_range, cv_accuracies, marker='s', linestyle='--', color='green', label='CV Accuracy')
plt.title('Accuracy of different k values')
plt.xlabel('k value')
plt.ylabel('Accuracy')
plt.grid(True)
plt.legend()  # 显示图例

# 绘制垂线标记最优k值
plt.axvline(x=best_k, color='red', linestyle='-')

# 标记最优k值和对应的准确率
best_test_accuracy = test_accuracies[best_k - 1]  # k从1开始，索引需要减1
plt.text(best_k + 0.5, best_test_accuracy, 
         f'k={best_k}, Test Accuracy={best_test_accuracy:.3f}', 
         color='red')

# 保存为PDF文件
plt.savefig('accuracy_plot.pdf')
plt.close()
print("准确率折线图已保存为 accuracy_plot.pdf")

# 将最佳KNN模型保存到二进制文件
with open("best_knn_model.pkl", "wb") as f:
    pickle.dump(best_knn, f)

print("最优KNN模型已保存为 best_knn_model.pkl")